from dotenv import load_dotenv
from functools import wraps
from flask import Flask, jsonify, Response, request, redirect, url_for
import flask
import os
from cache import MemoryCache
import cx_Oracle
from load_ddl import connect_to_oracle, get_table_metadata, get_column_metadata, generate_ddl_statements, train_vanna_with_ddl, generate_example_queries, train_vanna_with_queries

app = Flask(__name__, static_url_path='')

# SETUP
cache = MemoryCache()

from vanna.remote import VannaDefault
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
from vanna.flask import VannaFlaskApp
from openai import OpenAI
import pandas as pd
import pymysql

client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    default_headers = {"x-foo": "true"}
)

class MyVanna(ChromaDB_VectorStore,OpenAI_Chat):
    def __init__(self,client=None,config=None):
        ChromaDB_VectorStore.__init__(self,config=config)
        OpenAI_Chat.__init__(self,client=client,config=config)

vn = MyVanna(client=client,config={"model": "qwen2.5-7b-instruct",})   #chatglm2-6b
vn.max_tokens = 800
vn.temperature = 0.5

# 数据库连接配置
config = {
    'username': 'RESM_WHJT_PROD',
    'password': 'rZzj82q5Q2QBZ765zR6sU',
    'host': '192.168.2.105',
    'port': '1521',
    'service_name': 'ORCLPDB1'
}

# 连接数据库并训练
try:
    # 连接数据库
    connection = connect_to_oracle(**config)
    
    # 获取元数据
    tables_df = get_table_metadata(connection)
    columns_df = get_column_metadata(connection)
    
    print(f"获取到 {len(tables_df)} 个表的元数据")
    
    # 1. 首先训练 DDL
    print("开始训练表结构(DDL)...")
    ddl_list = generate_ddl_statements(tables_df, columns_df)
    print(f"生成了 {len(ddl_list)} 个 DDL 语句")
    ddl_training_ids = train_vanna_with_ddl(vn, ddl_list)
    print(f"DDL训练完成，共训练 {len(ddl_training_ids)} 个表结构")
    
    # 2. 然后训练示例查询
    print("\n开始训练示例查询...")
    query_list = generate_example_queries(tables_df, columns_df)
    print(f"生成了 {len(query_list)} 个查询语句")
    query_training_ids = train_vanna_with_queries(vn, query_list)
    print(f"查询训练完成，共训练 {len(query_training_ids)} 个查询")
    
    # 获取训练数据总数
    training_data = vn.get_training_data()
    print("\n训练总结:")
    print(f"成功训练 {len(ddl_training_ids)} 个表结构")
    print(f"成功训练 {len(query_training_ids)} 个示例查询")
    print(f"Milvus 中总共有 {len(training_data)} 条训练数据")

except cx_Oracle.Error as error:
    print(f"数据库错误: {error}")
except Exception as e:
    print(f"其他错误: {str(e)}")
    import traceback
    traceback.print_exc()
finally:
    if 'connection' in locals():
        connection.close()

vn.connect_to_oracle(dsn='192.168.2.105:1521/ORCLPDB1', user='RESM_WHJT_PROD', password='rZzj82q5Q2QBZ765zR6sU')
